import os
import shutil
import argparse
from pathlib import Path


def create_dir_if_not_exists(path):
    """创建目录（如果不存在）"""
    if not os.path.exists(path):
        os.makedirs(path)


def is_split_format(directory):
    """
    判断目录是否为按类别分文件夹格式
    如果目录中大多数条目是子目录，则认为是split格式
    """
    items = os.listdir(directory)
    if not items:
        return False  # 空目录

    # 统计子目录数量
    dir_count = 0
    file_count = 0

    for item in items:
        item_path = os.path.join(directory, item)
        if os.path.isdir(item_path):
            dir_count += 1
        else:
            file_count += 1

    # 如果子目录数量多于文件，认为是split格式
    return dir_count > file_count


def has_class_prefix(filename):
    """判断文件名是否包含类别前缀（格式: class_xxx）"""
    return '_' in filename and len(filename.split('_', 1)) >= 2


def is_flat_format(directory):
    """
    判断目录是否为单文件夹含类别前缀格式
    如果目录中大多数文件有类别前缀，则认为是flat格式
    """
    items = os.listdir(directory)
    if not items:
        return False  # 空目录

    # 统计带有类别前缀的文件数量
    prefix_files = 0
    total_files = 0

    for item in items:
        item_path = os.path.join(directory, item)
        if os.path.isfile(item_path):
            total_files += 1
            if has_class_prefix(item):
                prefix_files += 1

    # 如果超过一半的文件有类别前缀，认为是flat格式
    return total_files > 0 and (prefix_files / total_files) > 0.5


def detect_format(directory):
    """检测目录的格式"""
    if not os.path.isdir(directory):
        raise ValueError(f"{directory} 不是一个有效的目录")

    is_split = is_split_format(directory)
    is_flat = is_flat_format(directory)

    if is_split and not is_flat:
        return "split"
    elif is_flat and not is_split:
        return "flat"
    elif is_split and is_flat:
        # 混合格式，优先判断为split
        return "split"
    else:
        # 无法明确判断的格式
        return None


def split_to_classes(src_dir, dest_dir):
    """
    将单文件夹多文件格式转换为按类别分文件夹格式
    输入格式: class1_1.png, class1_2.png, class2_1.png...
    输出格式:
        class1/
            1.png, 2.png...
        class2/
            1.png, 2.png...
    """
    create_dir_if_not_exists(dest_dir)

    # 遍历源目录中的所有文件
    for filename in os.listdir(src_dir):
        file_path = os.path.join(src_dir, filename)

        # 跳过目录，只处理文件
        if os.path.isdir(file_path):
            continue

        # 分割文件名获取类别
        parts = filename.split('_', 1)  # 只分割一次
        if len(parts) < 2:
            print(f"警告: 文件名 {filename} 不符合格式，跳过")
            continue

        class_name, rest = parts
        class_dir = os.path.join(dest_dir, class_name)
        create_dir_if_not_exists(class_dir)

        # 构建目标路径
        dest_path = os.path.join(class_dir, rest)

        # 复制文件
        shutil.copy2(file_path, dest_path)
        # print(f"复制 {filename} 到 {class_dir}")


def merge_to_flat(src_dir, dest_dir):
    """
    将按类别分文件夹格式转换为单文件夹多文件格式
    输入格式:
        class1/
            1.png, 2.png...
        class2/
            1.png, 2.png...
    输出格式: class1_1.png, class1_2.png, class2_1.png...
    """
    create_dir_if_not_exists(dest_dir)

    # 遍历源目录中的所有类别文件夹
    for class_name in os.listdir(src_dir):
        class_dir = os.path.join(src_dir, class_name)

        # 跳过非目录文件
        if not os.path.isdir(class_dir):
            continue

        # 遍历类别文件夹中的所有文件
        for filename in os.listdir(class_dir):
            file_path = os.path.join(class_dir, filename)

            # 跳过子目录
            if os.path.isdir(file_path):
                print(f"警告: {class_dir} 中包含子目录 {filename}，跳过")
                continue

            # 构建新文件名
            new_filename = f"{class_name}_{filename}"
            dest_path = os.path.join(dest_dir, new_filename)

            # 处理文件名冲突
            counter = 1
            while os.path.exists(dest_path):
                name, ext = os.path.splitext(filename)
                new_filename = f"{class_name}_{name}_{counter}{ext}"
                dest_path = os.path.join(dest_dir, new_filename)
                counter += 1

            # 复制文件
            shutil.copy2(file_path, dest_path)
            # print(f"生成 {new_filename}")


def process_with_train_val(src_root, dest_root):
    """
    处理包含train/val子目录的数据集转换
    保持train和val的划分结构不变
    """
    # 定义训练集和验证集目录名
    splits = ['train', 'val']
    mappings_generated = False  # 标记是否已生成映射（确保train/val共享同一映射）

    for split in splits:
        src_split_dir = os.path.join(src_root, split)
        if not os.path.exists(src_split_dir):
            raise ValueError(f"错误: 未找到必需的 {split} 目录，请检查源目录结构")

        # 检测当前split目录的格式
        split_format = detect_format(src_split_dir)
        if split_format is None:
            raise ValueError(f"错误: 无法识别 {split} 目录的格式，请检查目录结构")

        # 创建目标split目录
        dest_split_dir = os.path.join(dest_root, split)
        create_dir_if_not_exists(dest_split_dir)

        # 执行转换
        print(f"处理 {split} 集，格式: {split_format}")
        if split_format == 'split':
            merge_to_flat(src_split_dir, dest_split_dir)
            if not mappings_generated:
                generate_class_mappings(mapping_dir=dest_root, source_dir=dest_split_dir)
                mappings_generated = True
        else:  # flat格式
            split_to_classes(src_split_dir, dest_split_dir)


def generate_class_mappings(classes=None, mapping_dir=None, source_dir=None):
    """
    生成并保存类别与索引之间的映射关系

    参数:
        classes: 类别名称的可迭代对象，若为None则从source_dir提取
        mapping_dir: 映射文件保存的目录路径
        source_dir: 当classes为None时，用于提取类别的源文件目录

    返回:
        classes_to_idx: 类别到索引的映射字典
        idx_to_labels: 索引到类别的映射字典
    """
    # 校验参数合法性
    if classes is None:
        if source_dir is None or not os.path.exists(source_dir):
            raise ValueError("当classes为None时，必须提供有效的source_dir用于提取类别")

        # 从源目录提取类别（沿用原逻辑：按文件名前缀分割）
        all_classes = set()
        for filename in os.listdir(source_dir):
            # 跳过非文件类型（如子目录）
            if not os.path.isfile(os.path.join(source_dir, filename)):
                continue
            # 提取类别（假设文件名格式为'class_数字_...'）
            class_name = filename.split('_')[0]
            all_classes.add(class_name)
        classes = all_classes

    if mapping_dir is None:
        raise ValueError("必须指定映射文件保存目录mapping_dir")

    # 确保保存目录存在
    os.makedirs(mapping_dir, exist_ok=True)

    # 对类别进行排序并创建映射
    sorted_classes = sorted(classes)
    classes_to_idx = {cls: i for i, cls in enumerate(sorted_classes)}
    idx_to_labels = {i: cls for i, cls in enumerate(sorted_classes)}

    # 保存classes_to_idx映射
    with open(os.path.join(mapping_dir, 'classes_to_label.txt'), 'w', encoding='utf-8') as f:
        for cls, idx in classes_to_idx.items():
            f.write(f"{cls},{idx}\n")

    # 保存idx_to_labels映射
    with open(os.path.join(mapping_dir, 'label_to_classes.txt'), 'w', encoding='utf-8') as f:
        for idx, cls in idx_to_labels.items():
            f.write(f"{idx},{cls}\n")

    # 保存类别列表
    with open(os.path.join(mapping_dir, 'classes.txt'), 'w', encoding='utf-8') as f:
        for cls in sorted_classes:
            f.write(f"{cls}\n")

    print(f"类别映射已保存至 {mapping_dir}")
    print(f"类别列表已保存至 {os.path.join(mapping_dir, 'classes.txt')}")

def main():
    parser = argparse.ArgumentParser(description='处理包含train和val子目录的数据集格式转换工具')
    parser.add_argument('--source', default=r'D:\Code\2-ZSL\Zero-Shot-Learning\data\split\PU', help='源数据集根目录（如3DChairs）')
    parser.add_argument('--dest', default=r'D:\Code\2-ZSL\Zero-Shot-Learning\data\split\PU-原始图片', help='目标数据集根目录')

    args = parser.parse_args()

    # 验证源目录存在
    if not os.path.isdir(args.source):
        raise ValueError(f"错误: 源目录 {args.source} 不存在")

    print(f"处理包含train/val划分的数据集: {args.source}")
    process_with_train_val(args.source, args.dest)

    print("转换完成!")


if __name__ == "__main__":
    main()
